import os,sys
import re
from textgrid import TextGrid, IntervalTier
import numpy as np

class TextGrid_diff_analysis:
    def __init__(self, srctg, dsttg, src_tiers=0, dst_tiers=0):
        self._addcount = 0
        self._subcount = 0
        self._errcount = 0
        self._samecount = 0
        self._allcount = 0

        self._cer = 0.0

        self._diffstring = ''
        
        self._src_name = os.path.basename(srctg)
        self._dst_name = os.path.basename(dsttg)

        self._srctg = TextGrid()
        self._dsttg = TextGrid()

        self._src_tiers = src_tiers
        self._dst_tiers = dst_tiers

        self._srctg.read(srctg)
        self._dsttg.read(dsttg)


    def equals(self, char_a, char_b):
        if(str.lower(char_a) == str.lower(char_b)):
            return True
        # if str.lower(interval_a.mark) == str.lower(interval_b.mark):
        #     return True
        return False


    def TextGrid_to_text(self):
        seq = []
        for i in range(len(self._srctg.tiers[self._src_tiers])):
            if self._srctg.tiers[self._src_tiers][i].mark != '':
                seq.append(self._srctg.tiers[self._src_tiers][i].mark)
        
        # 去除非音素标记，不同数据集标注可能不同
        # 此时针对BBSpeech
        seq_nonsil = []
        for i in range(len(self._srctg.tiers[self._src_tiers])):
            if self._srctg.tiers[self._src_tiers][i].mark != 'sil':
                seq_nonsil.append(self._srctg.tiers[self._src_tiers][i].mark)

        seq1 = []
        for i in range(len(self._dsttg.tiers[self._dst_tiers])):
            # 对于mfa，去除识别头部的sil，以评估准确的纯音素结果
            if self._dsttg.tiers[self._dst_tiers][i].mark != '' and self._dsttg.tiers[self._dst_tiers][i].mark != 'sil':
                seq1.append(self._dsttg.tiers[self._dst_tiers][i].mark)
            # 对齐结果中无法输出sil，相应mark皆为空，造成计算时识别序列错位；暂未解决，使用#占位
            else:
                seq1.append('#')
        
        return seq, seq_nonsil, seq1


    def lcsc(self, seqx, seqy):
        lenx = len(seqx)
        leny = len(seqy)

        table = [[[] for x in range(leny + 1)] for y in range(lenx + 1)]
        Matrix = [[0 for x in range(leny + 1)] for y in range(lenx + 1)]
        for tmp in range(0, lenx + 1):
            Matrix[tmp][0] = tmp
        for tmp in range(0, leny + 1):
            Matrix[0][tmp] = tmp
        
        Matrix[0][0] = 0

        for xline in range(1, lenx + 1):
            for yline in range(1, leny + 1):
                MinCost1 = Matrix[xline - 1][yline] + 1 # src delete
                MinCost2 = Matrix[xline][yline - 1] + 1 # src add
                MinCost = min(MinCost1, MinCost2)

                if self.equals(seqx[xline - 1], seqy[yline - 1]):
                    ReplaceCost = 0
                else:
                    ReplaceCost = 1 # src replace
                
                if ReplaceCost + Matrix[xline - 1][yline - 1] < MinCost:
                    MinCost =  ReplaceCost + Matrix[xline - 1][yline - 1]
                    table[xline][yline].extend(table[xline - 1][yline - 1])
                    if ReplaceCost == 0:
                        table[xline][yline].append([xline-1, yline-1])
                elif MinCost2 == MinCost1:
                    if len(table[xline][yline - 1]) >= len(table[xline - 1][yline]):
                        table[xline][yline] = table[xline][yline - 1]
                    else:
                        table[xline][yline] = table[xline - 1][yline]
                elif MinCost2 < MinCost1:
                    table[xline][yline] = table[xline][yline - 1]
                else:
                    table[xline][yline] = table[xline - 1][yline]

                Matrix[xline][yline] = MinCost
        
        seqcls = table[lenx][leny]
        seqcls.append([lenx, leny])
        return seqcls


    def MatchSeq_to_TimeSeq(self, matchseq):
        srclen = len(self._srctg[self._src_tiers])
        dstlen = len(self._dsttg[self._dst_tiers])
        src_time = []
        dst_time = []
        diff_time = []

        for i in matchseq:
            if (i[0] < srclen) and (i[1] < dstlen):
                if self._srctg[self._src_tiers][i[0]].mark != self._dsttg[self._dst_tiers][i[1]].mark:
                    print('Something Wrong in matchseq!') # 应当输出出错文件号，匹配序列
                    sys.exit()

                srcMintime = self._srctg[self._src_tiers][i[0]].minTime
                srcMaxtime = self._srctg[self._src_tiers][i[0]].maxTime
                dstMintime = self._dsttg[self._dst_tiers][i[1]].minTime
                dstMaxtime = self._dsttg[self._dst_tiers][i[1]].maxTime

                src_time.append((srcMintime, srcMaxtime))
                dst_time.append((dstMintime, dstMaxtime))

                diff = abs(srcMintime - dstMintime) + abs(srcMaxtime - dstMaxtime)
                diff_time.append(diff * 1000)

        return src_time, dst_time, diff_time

    
    def MatchSeq_to_TextSeq(self, matchseq, srcseq, dstseq):
        if len(dstseq) <= 0:
            self._allcount = len(srcseq)
            if self._allcount == 0:
                self._cer = 0.0
                return 0
            self._diffstring = '[ ' + ' '.join(srcseq) + ' ]'
            self._subcount = self._allcount
            self._cer = 1.0
            return 0
    
        lcsparts = matchseq

        if len(lcsparts) == 0:
            self._errcount = len(dstseq)
            self._allcount = len(srcseq)
            self._cer = 1.0 * (self._errcount) / self._allcount
            return 0
        
        allcount = len(srcseq)
        samecount = 0

        diffparts = []

        addcount = 0
        subcount = 0
        errcount = 0

        srcindex = 0
        dstindex = 0
        lcsindex = 0

        while lcsindex < len(lcsparts):
            src = lcsparts[lcsindex][0]
            dst = lcsparts[lcsindex][1]

            while srcindex < src and dstindex < dst:
                diffparts.append('(')
                diffparts.append(srcseq[srcindex])
                diffparts.append(':')
                diffparts.append(dstseq[dstindex])
                diffparts.append(')')

                errcount = errcount + 1

                srcindex = srcindex + 1
                dstindex = dstindex + 1
            
            if srcindex < src:
                diffparts.append('[')
                while srcindex < src:
                    diffparts.append(srcseq[srcindex])

                    subcount = subcount + 1
                    srcindex = srcindex + 1
                diffparts.append(']')

            if dstindex < dst:
                diffparts.append('<')
                while dstindex < dst:
                    diffparts.append(dstseq[dstindex])

                    addcount = addcount + 1
                    dstindex = dstindex + 1
                diffparts.append('>')

            if lcsindex != len(lcsparts) - 1:
                diffparts.append(dstseq[dstindex])

                samecount = samecount + 1

                srcindex = srcindex + 1
                dstindex = dstindex + 1
            lcsindex = lcsindex + 1

        self._diffstring = ''
        for part in diffparts:
            self._diffstring = self._diffstring + ' ' + part

        try:
            cer = 1.0 * (addcount + subcount + errcount) / len(srcseq)
        except ZeroDivisionError:
            if addcount + subcount + errcount == 0:
                cer = 0.0
            else:
                cer = float('inf')

        self._cer = cer
        self._addcount = addcount
        self._subcount = subcount
        self._errcount = errcount

        self._allcount = allcount
        self._samecount = samecount
        if samecount != len(lcsparts) - 1:
            print ('error samecount is not equal lcsparts length')
            sys.exit()

            return 0

        return len(diffparts)


def plot_diff(diffT, set_name, out_dir):
    import matplotlib.pyplot as plt
    import matplotlib.ticker as ticker
    import seaborn as sns
    import statsmodels

    sns.set()
    fig, axes = plt.subplots(1, 2, figsize = (20,10))
    # sns.distplot(diffT, rug=True, ax=axes[0])
    sns.kdeplot(diffT, cumulative=True, ax=axes[0])
    sns.distplot(diffT, rug=True, kde=False, ax=axes[1])
    
    axes[0].set_xlabel("Abs_diff/ms")
    axes[0].set_ylabel("CDF")
    axes[0].xaxis.set_major_locator(ticker.MultipleLocator(10))
    axes[0].yaxis.set_major_locator(ticker.MultipleLocator(0.1))

    axes[1].set_xlabel("Abs_diff/ms")
    axes[1].set_ylabel("Frequency")
    fig.savefig(f'{out_dir}/{set_name}.sns.png')


def process(set_name, labeled_file, recog_file, result_dir):
    labeled_set = {}
    recog_set = {}
    diffTime = []

    total_add, total_del, total_err, total_cost, total_len = 0, 0, 0, 0, 0
    err_sentence, total_sentence = 0, 0

    stats = {}

    # 记录大差距文件
    large_diff = {}
    total_large = 0
    num_200 = 0
    num_500 = 0
    num_fail = 0

    for i in open(labeled_file):
        tmp = i.strip().split()
        labeled_set[tmp[0]] = tmp[1]
    for i in open(recog_file):
        tmp = i.strip().split()
        recog_set[tmp[0]] = tmp[1]
    for i in labeled_set.keys():
        if i not in recog_set:
            print(f'Cannot find {i} in {recog_file}, check!')
            continue

        # src_tiers 的传输！！！
        textgrid_diff = TextGrid_diff_analysis(labeled_set[i], recog_set[i], src_tiers=0, dst_tiers=0)
        src_string, seq_nonsil, dst_string = textgrid_diff.TextGrid_to_text()
        matchseq = textgrid_diff.lcsc(src_string, dst_string)
        src_time, dst_time, diff_time = textgrid_diff.MatchSeq_to_TimeSeq(matchseq)

        # 可以考虑不对所有项进行此步
        _ = textgrid_diff.MatchSeq_to_TextSeq(matchseq,src_string, dst_string)

        # 常规信息
        total_sentence += 1
        more = textgrid_diff._addcount
        less = textgrid_diff._subcount
        err = textgrid_diff._errcount
        cost = more + less + err
        err_rate = textgrid_diff._cer
        flag = textgrid_diff._diffstring

        if err_rate != 0:
            err_sentence += 1
        err_show = "total: %d, add: %d, del: %d, err: %d, cer: %.4f%% " % (len(src_string), more, less, err, err_rate*100)
        stats[i.strip()] = i.strip() + "\t" + ' '.join(src_string) + "\n" + \
                           i.strip() + "\t" + ' '.join(dst_string) + "\n" + \
                           "Nonsilence_ph:\t" + ' '.join(seq_nonsil) + "\n" + \
                           flag.strip() + "\n" + \
                           i.strip() + "\t" + str(src_time) + "\n" + \
                           i.strip() + "\t" + str(dst_time) + "\n" + \
                           "Absolute_diff:\t" + str(diff_time) + "\n" + \
                           str(f'mean: {np.mean(diff_time)} ; median:{np.median(diff_time)}\n') + \
                           err_show + "\n\n"

        total_add += more
        total_del += less
        total_err += err
        total_cost += cost
        total_len += len(src_string)

        # 大差距信息
        if max(diff_time) > 100:

            large_diff[i.strip()] = i.strip() + "\t" + ' '.join(src_string) + "\n" + \
                           i.strip() + "\t" + ' '.join(dst_string) + "\n" + \
                           "Nonsilence_ph:\t" + ' '.join(seq_nonsil) + "\n" + \
                           flag.strip() + "\n" + \
                           i.strip() + "\t" + str(src_time) + "\n" + \
                           i.strip() + "\t" + str(dst_time) + "\n" + \
                           "Absolute_diff:\t" + str(diff_time) + "\n" + \
                           str(f'mean: {np.mean(diff_time)} ; median:{np.median(diff_time)}\n') + \
                           err_show + "\n\n"

            for t in diff_time:
                if t <= 100:
                    diffTime.append(t)
                elif t <= 200:
                    num_200 += 1
                elif t <= 500:
                    num_500 += 1
                else:
                    num_fail += 1
        else:
            diffTime.extend(diff_time)

    np.savez_compressed(f'{result_dir}/diffTime', diffTime)
    plot_diff(diffTime, set_name, result_dir)

    try:
        total_err_rate = total_cost / total_len * 100
    except ZeroDivisionError:
        if total_cost == 0:
            total_err_rate = 0.0
        else:
            total_err_rate = float('inf')
    try:
        toal_err_rate_exclude_add = (total_del + total_err) / total_len * 100
    except ZeroDivisionError:
        if total_del + total_err != 0:
            toal_err_rate_exclude_add = float('inf')
        else:
            toal_err_rate_exclude_add = 0.0
    final = "total: %d, add: %d, del: %d, err: %d, cer: %.4f%%, ser: %.2f%%, cer(exclude_add): %.4f%%" % (total_len, total_add, total_del, total_err, total_err_rate, err_sentence / total_sentence *100, toal_err_rate_exclude_add)
    print(final)

    final_diffT = "total_recogPh: %d, <100_recogPh: %d, <200_recogPh: %d, <500_recogPh: %d, failed_recogPh: %d, mean_AbsDiffTime(ms): %.3f, madian_AbsDiffTime(ms): %.3f" % (total_len-total_del-total_err, len(diffTime), num_200, num_500, num_fail, np.mean(diffTime), np.median(diffTime))
    print(final_diffT)

    with open(f'{result_dir}/result_diff.txt', 'w', encoding="utf-8") as f:
        f.write(f'file:{result_dir}/result_diff.txt\n\n')
        f.write("说明：[]——丢字，()——错字，<>——多字，cer——字错误率，ser——句错误率\n\t仅考虑标注文本中的中英字符及阿拉伯数字\n\n")
        f.write("=====================================\n")
        for i in sorted(stats.keys()):
            f.write(stats[i])
        f.write("=====================================\n")
        f.write("STAT：" + final + '\n')
        f.write("PhTime_AbsDiff: " + final_diffT + '\n')

    with open(f'{result_dir}/large_diff.txt', 'w', encoding="utf-8") as f:
        for i in sorted(large_diff.keys()):
            f.write(large_diff[i])
        f.write(str(total_large) + '\n')


    
if __name__ == '__main__':
    if len(sys.argv) != 5:
        print ("Usage: python3 count_diff_textgrid.py set_name, labeled_file recognized_file result_dir")
        sys.exit(1)

    set_name = sys.argv[1]
    labeled_file = sys.argv[2]
    recognized_file = sys.argv[3]
    result_dir = sys.argv[4]

    os.makedirs(result_dir, exist_ok=True)
    process(set_name, labeled_file, recognized_file, result_dir)
